package com.example.xiaosen.controller;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.example.xiaosen.entity.vo.Result;
import com.example.xiaosen.repository.ChatHistoryRepository;
import com.example.xiaosen.repository.FileRepository;
import com.example.xiaosen.service.IResourceService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.transformer.KeywordMetadataEnricher;
import org.springframework.ai.model.transformer.SummaryMetadataEnricher;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.join.DocumentJoiner;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @Author: zhouxj@ferrotec.com.cn
 * @version: v1.0
 * @Date: 2025/5/14 13:14
 * @Description:
 */
@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {
    private final FileRepository fileRepository;
    private final VectorStore vectorStore;
    private final ChatHistoryRepository chatRepository;
    private final ChatClient pdfChatClient;
    private final OpenAiChatModel chatModel;
    @Autowired
    private IResourceService resourceService;
    /**
     * 文件上传，知识库存储
     * SpringAI支持的文件读取器如下：
     * {@link org.springframework.ai.document.DocumentReader}
     *  TikaDocumentReader：apache tika内容分析工具包，它能够检测和提取超过一千种不同文件类型（如 PPT、XLS 和 PDF）的元数据和文本。
     *  JsonReader： 读取json文件
     *  ParagraphPdfDocumentReader： 将PDF拆分成段落，每个段落作为一个Document,需要pdf设置大纲
     *  TextReader：读取文本文件
     *  PagePdfDocumentReader ：将PDF拆分成每一页作为一个Document
     */
    @RequestMapping("/upload/{chatId}")
    public Result uploadPdf(@PathVariable String chatId,
                            @RequestParam("file") MultipartFile file) {
        try {
            // 校验文件格式
            if(!MediaType.APPLICATION_PDF.toString().equals(file.getContentType())){
                return Result.fail("文件格式错误");
            }
            // 保存文件
            Resource resource = file.getResource();
            if (!fileRepository.save(chatId, resource)) {
                return Result.fail("文件保存失败");
            }
            // 使用apache tika读取pdf内容
            TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
            List<Document> tikaRead = tikaDocumentReader.read();
            // TokenTextSplitter用于切割Document文档，可以设置分块大小
            List<Document> documents1 = new TokenTextSplitter().apply(tikaRead);
            // 添加权限元数据
            documents1.stream().forEach(document -> {
                document.getMetadata().put("role","user");
            });
            // 保存文本向量，springAI会先调用嵌入模型将文档向量化，然后保存到向量数据库中
            // 关键词提取添加到元数据中
            KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(chatModel, 5);
            // 摘要总结提取添加到元数据中
            SummaryMetadataEnricher summaryMetadataEnricher = new SummaryMetadataEnricher(
                    chatModel,
                    List.of(SummaryMetadataEnricher.SummaryType.CURRENT, SummaryMetadataEnricher.SummaryType.NEXT)
            );
            List<Document> enrichedDocs1 = keywordMetadataEnricher.apply(documents1);
            List<Document> enrichedDocs2 = summaryMetadataEnricher.apply(enrichedDocs1);
            vectorStore.add(enrichedDocs2);

            // 直接使用spring ai提供的pdf读取器，将pdf拆分成每一页作为一个Document
//            PagePdfDocumentReader reader = new PagePdfDocumentReader(resource,
//                    PdfDocumentReaderConfig.builder()
//                            .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
//                            // 每1页PDF作为一个Document
//                            .withPagesPerDocument(1)
//                            .build());
//            List<Document> documents = reader.read();
//            vectorStore.add(documents);
            // metadata条件删除
//            FilterExpressionBuilder b = new FilterExpressionBuilder();
//            vectorStore.delete(b.eq("source", "hello").build());
            return Result.ok();
        } catch (Exception e) {
            return Result.fail(e.getMessage());
        }
    }

    /**
     * 文件下载
     */
    @GetMapping("/file/{chatId}")
    public ResponseEntity<Resource> download(@PathVariable("chatId") String
                                                     chatId) throws IOException {
        // 1.读取文件
        Resource resource = fileRepository.getFile(chatId);
        if (!resource.exists()) {
            return ResponseEntity.notFound().build();
        }
        // 2.文件名编码，写入响应头
        String filename =
                URLEncoder.encode(Objects.requireNonNull(resource.getFilename()),
                        StandardCharsets.UTF_8);
        // 3.返回文件
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header("Content-Disposition", "attachment; filename=\"" +
                        filename + "\"")
                .body(resource);
    }


    /**
     * 大模型基于RAG的pdf问答助手
     * @param prompt
     * @param chatId
     * @return
     */
    @RequestMapping(value = "/chat", produces = "text/html;charset=UTF-8")
    public Flux<String> chat(String prompt, String chatId) {
        chatRepository.save("pdf", chatId);
        // 这里设计为一个会话对应一个pdf，根据文件名进行向量数据库过滤
        String fileName = resourceService.getOne(new LambdaQueryWrapper<com.example.xiaosen.entity.po.Resource>()
                .eq(StringUtils.isNotEmpty(chatId), com.example.xiaosen.entity.po.Resource::getChatId, chatId)).getFileName();
        PromptTemplate emptyContextPromptTemplate = new PromptTemplate("""
                你应该输出下面的内容：
                抱歉，我只能回答恋爱相关的问题，别的没办法帮到您哦，
                有问题可以联系编程导航客服 https://codefather.cn
                """);
        // RetrievalArgumentAdvisor配置
        // 检索前优化→多源检索→检索后处理
        RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
                // 检索前阶段，查询转换器，设置一系列transformedQuery对原始用户查询进行转换
                .queryTransformers(
                        // 利用大语言模型重写用户查询
                        RewriteQueryTransformer.builder()
                                .chatClientBuilder(pdfChatClient.mutate())
                                .build(),
                        // 压缩上下文，避免token超限
                        CompressionQueryTransformer.builder()
                                .chatClientBuilder(pdfChatClient.mutate())
                                .build())
                // 检索前阶段，查询扩展器，通过大语言模型将用户原始查询进行优化，生成3-5条变体
                .queryExpander(MultiQueryExpander.builder()
                        .chatClientBuilder(pdfChatClient.mutate())
                        .numberOfQueries(3)
                        // 包含原始查询
                        .includeOriginal(true)
                        .build())
                // 检索阶段，文档检索器，负责从底层数据源（例如搜索引擎、矢量存储、 数据库或知识图谱）检索返回List<Document>
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(vectorStore)
                        .similarityThreshold(0.5)
                        .topK(5)
                        .filterExpression(new FilterExpressionBuilder()
                                .eq("role", "user")
                                .build())
                        .build())
                // 检索后阶段，查询增强器，默认情况下，不允许检索到的上下文为空。当这种情况发生时，它指示模型不回答用户查询。可以按如下方式配置空上下文处理方式：
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(false)
                        .emptyContextPromptTemplate(emptyContextPromptTemplate)
                        .build())
                .build();
        // 文档联结，该组件将基于多查询和多数据源检索的文档合并为单一文档集合，在联结过程中可处理重复文档及互惠排名策略。
//        Map<Query, List<List<Document>>> documentsForQuery = ...
//        DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
//        List<Document> documents = documentJoiner.join(documentsForQuery);
        return pdfChatClient
                .prompt(prompt)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                // 对QuestionAnswerAdvisor设置过滤规则，确保会话只应用自己的知识库，即与pdf文件一一对应
                // 这里匹配的是向量数据库中存储的metadata元数据内容，每种文件读取器读取的元数据格式不同，如PagePdfDocumentReader中metadata为（file_name、page_number）
                // 而TikaDocumentReader则对应文件名的元数据为source
                .advisors(ragAdvisor)
                .stream()
                .content();
    }
}